#!/usr/bin/env python

import roslib
roslib.load_manifest('ipm_to_laserscan')
import sys
import rospy
import cv
import yaml
from std_msgs.msg import String
from sensor_msgs.msg import Image
from cv_bridge import CvBridge, CvBridgeError

class ipm_node:
    CHECKERS_WIDE = 7
    CHECKERS_TALL = 9
    CHECKER_WIDTH_IN_M = 0.065

    CAMERA_OFFSET = 2.0

    VIEW_HEIGHT = 4.0
    VIEW_WIDTH = 4.0

    IN_TOPIC = "/camera/image_raw"
    OUT_TOPIC = "/camera/ipm"

    def __init__(self):
        rospy.init_node('ipm_node')
        self.image_in_topic = rospy.get_param('~image_in_topic', self.IN_TOPIC)
        self.image_out_topic = rospy.get_param('~image_out_topic', self.OUT_TOPIC)
        self.checkers_wide = rospy.get_param('~checkers_wide', self.CHECKERS_WIDE)
        self.checkers_tall = rospy.get_param('~checkers_tall', self.CHECKERS_TALL)
        self.checker_width = rospy.get_param('~checker_width', self.CHECKER_WIDTH_IN_M)
        self.camera_offset = rospy.get_param('~camera_offset', self.CAMERA_OFFSET)
        self.view_h = rospy.get_param('~view_height', self.VIEW_HEIGHT)
        self.view_w = rospy.get_param('~view_width', self.VIEW_WIDTH)
        self.create_transform_map = rospy.get_param('~create_transform_map', False)
        self.map_file = rospy.get_param('~transform_file', None)

        self.show_window_param = rospy.get_param('~show_windows', True)
        self.load_transform_map = None
        self.dst = None

        if not self.create_transform_map:
            self.perspectiveMat = cv.CreateMat(3, 3, cv.CV_32FC1)

            f = open(self.map_file, 'r')
            config = yaml.load(f)

            rospy.loginfo("loading transform map %s"%self.map_file)

            real_pts = [(config['real_pts']['tl_x'],config['real_pts']['tl_y']),
                        (config['real_pts']['tr_x'],config['real_pts']['tr_y']),
                        (config['real_pts']['br_x'],config['real_pts']['br_y']),
                        (config['real_pts']['bl_x'],config['real_pts']['bl_y'])]

            img_pts = [(config['img_pts']['tl_x'],config['img_pts']['tl_y']),
                        (config['img_pts']['tr_x'],config['img_pts']['tr_y']),
                        (config['img_pts']['br_x'],config['img_pts']['br_y']),
                        (config['img_pts']['bl_x'],config['img_pts']['bl_y'])]

            self.resolution = config['resolution']
            self.ipm_image_w = config['img_w'] 
            self.ipm_image_h = config['img_h'] 

            cv.GetPerspectiveTransform(real_pts, img_pts, self.perspectiveMat)

            self.calibrated = True
        else:
            self.calibrated = False

        """ Give the OpenCV display window a name. """
        self.input_window = "Input Image"
        self.output_window = "Output Image"


        """ Create the window and make it re-sizeable (second parameter = 0) """
        if (self.show_window_param):
            cv.NamedWindow(self.input_window, 0)
            cv.NamedWindow(self.output_window, 0)

        """ Create the cv_bridge object """
        self.bridge = CvBridge()

        """ Subscribe to the raw camera image topic """
        self.image_sub = rospy.Subscriber(self.image_in_topic, Image, self.handle_image)

        """ Publish the IPM image to the image_ipm topic """
        self.image_pub = rospy.Publisher(self.image_out_topic, Image)


    def handle_image(self, data):

        try:
            """ Convert the raw image to OpenCV format """
            src = self.bridge.imgmsg_to_cv(data, "bgr8")
        except CvBridgeError, e:
            print e
            return
  
        (width, height) = cv.GetSize(src)

        if (not self.calibrated):
            print "calibrating..."
            """ find checkerboard corners in image 
                note: dimensions param asks for corners, not checker cells, hence the -1 """
            found_all, corners = cv.FindChessboardCorners( src, 
                            (self.checkers_wide-1, self.checkers_tall-1))

            pers_image = None

            if found_all:

                config = {}
                img_top_lt = (int(corners[0][0]), int(corners[0][1]))
                img_top_rt = (int(corners[self.checkers_wide-2][0]), 
                                int(corners[self.checkers_wide-2][1])) 
                img_btm_lt = (int(corners[(self.checkers_tall-2) * 
                                    (self.checkers_wide-1)][0]),
                              int(corners[(self.checkers_tall-2) * 
                                    (self.checkers_wide-1)][1])) 
                img_btm_rt = (int(corners[(self.checkers_tall-2) * 
                                (self.checkers_wide-1) + self.checkers_wide-2][0]),
                              int(corners[(self.checkers_tall-2) * 
                                (self.checkers_wide-1) + self.checkers_wide-2][1]))

                config['img_pts'] = {}
                config['img_pts']['tl_x'] = img_top_lt[0]
                config['img_pts']['tl_y'] = img_top_lt[1]
                config['img_pts']['tr_x'] = img_top_rt[0]
                config['img_pts']['tr_y'] = img_top_rt[1]
                config['img_pts']['bl_x'] = img_btm_lt[0]
                config['img_pts']['bl_y'] = img_btm_lt[1]
                config['img_pts']['br_x'] = img_btm_rt[0]
                config['img_pts']['br_y'] = img_btm_rt[1]

                img_pts = [ img_top_lt, 
                            img_top_rt,
                            img_btm_rt,
                            img_btm_lt ]

                # using the fact that there is less 
                # optical distortion closer to the cam
                # -- get pixel width of a checker square
                chkr_w = int((img_btm_rt[0] - img_btm_lt[0])/self.checkers_wide-2)
                self.rect_w = chkr_w * self.checkers_wide - 2
                self.rect_h = chkr_w * self.checkers_tall - 2
            
                self.resolution = self.checker_width / chkr_w
                print "Checker pixel size: %s"%chkr_w
                print "Pixel resolution (in meters): %s"%self.resolution
                rospy.loginfo("meters per px: %f"%self.resolution)

                config['resolution'] = self.resolution

                x_off = ((self.view_w / self.resolution) / 2) - (self.rect_w / 2)
                y_off = (self.view_h / self.resolution) - self.rect_h

                self.ipm_image_w = (self.view_w / self.resolution)
                self.ipm_image_h = (self.view_h / self.resolution)

                config['img_w'] = self.ipm_image_w
                config['img_h'] = self.ipm_image_h

                rospy.loginfo("image width in px: %f"%self.ipm_image_w)
                rospy.loginfo("xoff: %f"%x_off)
                rospy.loginfo("rect_w: %f"%self.rect_w)
                rospy.loginfo("rect_h: %f"%self.rect_h)

                real_top_lt = (x_off, y_off - chkr_w)
                real_top_rt = (x_off + self.rect_w, y_off - chkr_w)
                real_btm_rt = (x_off + self.rect_w, y_off + self.rect_h - chkr_w)
                real_btm_lt = (x_off, y_off + self.rect_h - chkr_w)

                config['real_pts'] = {}
                config['real_pts']['tl_x'] = real_top_lt[0]
                config['real_pts']['tl_y'] = real_top_lt[1]
                config['real_pts']['tr_x'] = real_top_rt[0]
                config['real_pts']['tr_y'] = real_top_rt[1]
                config['real_pts']['bl_x'] = real_btm_lt[0]
                config['real_pts']['bl_y'] = real_btm_lt[1]
                config['real_pts']['br_x'] = real_btm_rt[0]
                config['real_pts']['br_y'] = real_btm_rt[1]
                real_pts = [ real_top_lt, real_top_rt, real_btm_rt, real_btm_lt ]

                self.perspectiveMat = cv.CreateMat(3, 3, cv.CV_32FC1)


                cv.GetPerspectiveTransform(real_pts, img_pts, self.perspectiveMat)

                if (self.map_file is not None):
                    rospy.loginfo("saving transform map %s"%self.map_file)
                    f = open(self.map_file, 'w')
                    f.write(yaml.dump(config))

                self.dst = cv.CreateImage((self.ipm_image_w,
                                            self.ipm_image_h), cv.IPL_DEPTH_8U, 3)
                cv.WarpPerspective(src, self.dst, self.perspectiveMat, (cv.CV_INTER_LINEAR | 
                                                          cv.CV_WARP_INVERSE_MAP | 
                                                          cv.CV_WARP_FILL_OUTLIERS), 
                                                            cv.ScalarAll(0))

                """ throw some dots on the image to represent the found rect """
                cv.Circle(src, img_top_lt, 10, (255,0,0), -1)
                cv.Circle(src, img_top_rt, 10, (255,0,0), -1)
                cv.Circle(src, img_btm_rt, 10, (255,0,0), -1)
                cv.Circle(src, img_btm_lt, 10, (255,0,0), -1)

                cv.Circle(self.dst, real_top_lt, 10, (255,0,0),-1)
                cv.Circle(self.dst, real_top_rt, 10, (255,0,0),-1)
                cv.Circle(self.dst, real_btm_rt, 10, (255,0,0),-1)
                cv.Circle(self.dst, real_btm_lt, 10, (255,0,0),-1)

                cv.SaveImage("/tmp/calib_input.png", src)
                cv.SaveImage("/tmp/calib_output.png", self.dst)

                self.calibrated = True
            else:
                print "no calibration board found"
                return

        else:

            if (self.dst is None):
                self.dst = cv.CreateImage((self.ipm_image_w,self.ipm_image_h), cv.IPL_DEPTH_8U, 3)
                self.out = cv.CreateImage((self.ipm_image_w,self.ipm_image_h + self.camera_offset/self.resolution), cv.IPL_DEPTH_8U, 3)

            cv.WarpPerspective(src, self.dst, self.perspectiveMat, (cv.CV_INTER_LINEAR | 
                                                      cv.CV_WARP_INVERSE_MAP | 
                                                      cv.CV_WARP_FILL_OUTLIERS), 
                                                        cv.ScalarAll(0))

            cv.SetImageROI(self.out, (0,0, self.dst.width, self.dst.height))
            cv.Copy(self.dst,self.out)
            cv.ResetImageROI(self.out)

            if (self.show_window_param):
                cv.ShowImage(self.output_window, self.out)

            self.image_pub.publish(self.bridge.cv_to_imgmsg(self.out, "bgr8"))

        if (self.show_window_param):
            cv.ShowImage(self.input_window, src)

        cv.WaitKey(3)

def main(args):
      vn = ipm_node()
      try:
        rospy.spin()
      except KeyboardInterrupt:
        print "Shutting down vison node."
      cv.DestroyAllWindows()

if __name__ == '__main__':
    main(sys.argv)

